from pymilvus import MilvusClient, model

embedding_fn = model.DefaultEmbeddingFunction()

def get_embedding(text):

    vectors = embedding_fn.encode_documents(text)
    # The output vector has 768 dimensions, matching the collection that we just created.
    print("Dim:", embedding_fn.dim, vectors[0].shape)  # Dim: 768 (768,)

    # Each entity has id, vector representation, raw text, and a subject label that we use
    # to demo metadata filtering later.
    data = [
        {"id": i, "vector": vectors[i], "text": text[i], "subject": "history"}
        for i in range(len(vectors))
    ]
    print("Data has", len(data), "entities, each with fields: ", data[0].keys())
    print("Vector dim:", len(data[0]["vector"]))
    return data

client = MilvusClient("milvus_demo.db")

if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")

client.create_collection(
    collection_name="demo_collection",
    dimension=768,  # The vectors we will use in this demo has 768 dimensions
)

text = ['张三 18 岁了是一个男生', '李四 28 岁了是一个女生', '刘五 36 岁了是一个男生']

res = client.insert(collection_name="demo_collection", data=get_embedding(text))

query_vectors = embedding_fn.encode_queries(["女生"])

res = client.search(
    collection_name="demo_collection",  # target collection
    data=query_vectors,  # query vectors
    limit=10,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

print(res)

